import cv2
import matplotlib.pyplot as plt
import numpy as np

def unsharp_masking(image_path, gaussian_ksize=(5, 5), sigma=1.0, gain_factor=1.5, threshold=0):
    """
    非锐化掩膜（Unsharp Masking）的简化演示。

    参数:
    image_path (str): 输入图像的路径。
    gaussian_ksize (tuple): 高斯模糊的核大小，(宽, 高)。值越大，模糊程度越高。
                            必须是奇数，如 (5, 5)。
    gain_factor (float): 增益因子，决定细节叠加的强度。值越大，锐化效果越强。
                          通常在 1.0 到 3.0 之间。
    threshold (int): 阈值，小于此值的像素差值（细节）将被忽略，有助于抑制噪声。
                     通常在 0 到 10 之间。
    """

    # 1. 加载图像并转换为灰度图
    # 为了简化，我们只处理灰度图，彩色图需要对每个通道单独处理
    original = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)

    # 2. 生成模糊图像 (非锐化掩膜)
    # 使用高斯模糊来平滑图像，去除高频细节
    blurred = cv2.GaussianBlur(original, gaussian_ksize, sigma)
    print("步骤1: 已生成模糊图像。")

    # 3. 税化逻辑
    sharpened = float(gain_factor + 1) * original - float(gain_factor) * blurred
    sharpened = np.maximum(sharpened, np.zeros(sharpened.shape))
    sharpened = np.minimum(sharpened, 255 * np.ones(sharpened.shape))
    sharpened = sharpened.round().astype(np.uint8)

    if threshold > 0:
        low_contrast_mask = np.absolute(original - blurred) < threshold
        np.copyto(sharpened, original, where=low_contrast_mask)
        print(f"步骤3.5: 已应用阈值 {threshold} 抑制细节噪音。")

    # 4. 将强化的细节加回原图
    # 这里需要注意数据类型，避免减法溢出，所以先转为 float
    # 差值图像包含了高频信息（边缘和噪声）
    details = original.astype(np.float32) - blurred.astype(np.float32)
    print("步骤4: 已将细节加回原图，完成锐化。")

    # 5. 显示结果
    _plot_results(original, blurred, details, sharpened)


def _plot_results(original_image, blurred_image, detail_image, sharpened_image):
    # 为了可视化，将细节图像中心化并归一化到 0-255
    detail_image_viz = cv2.normalize(detail_image, None, 0, 255, cv2.NORM_MINMAX)
    detail_image_viz = detail_image_viz.astype(np.uint8)

    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 更好的中文字体，可以显示“-”号
    plt.figure(figsize=(15, 10))

    plt.subplot(2, 2, 1)
    plt.imshow(original_image, cmap='gray')
    plt.title('原始图像')
    plt.axis('off')

    plt.subplot(2, 2, 2)
    plt.imshow(blurred_image, cmap='gray')
    plt.title('模糊图像 (非锐化掩膜)')
    plt.axis('off')

    plt.subplot(2, 2, 3)
    plt.imshow(detail_image_viz, cmap='gray')
    plt.title('提取的细节图像 (已可视化)')
    plt.axis('off')

    plt.subplot(2, 2, 4)
    plt.imshow(sharpened_image, cmap='gray')
    plt.title('锐化后的图像')
    plt.axis('off')

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    # 请确保在运行前，当前目录下有一个名为 'input_image.jpg' 的图像文件
    # 你可以替换成你自己的图片路径
    # 也可以尝试不同的参数组合来观察效果
    unsharp_masking(
        image_path = './unsharp_masking/test_image.jpg',
        gaussian_ksize=(5, 5), # 模糊核大小：影响细节的粗细程度，越大越模糊，提取的细节越“粗”
        gain_factor=3,       # 增益因子：锐化强度，越大越锐利，也越容易引入副作用
        threshold=100            # 阈值：过滤掉微弱的细节（噪音），越大过滤越多
    )

    # 尝试不同参数组合
    # print("\n--- 尝试更强的锐化 ---")
    # unsharp_masking_demo('input_image.jpg', gaussian_ksize=(3, 3), gain_factor=2.5, threshold=2)

    # print("\n--- 尝试更柔和的锐化，更多降噪 ---")
    # unsharp_masking_demo('input_image.jpg', gaussian_ksize=(7, 7), gain_factor=1.0, threshold=10)